package test.tftp;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.SocketException;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import top.ieei.demo.tftp.Options;
import top.ieei.demo.tftp.TftpAck;
import top.ieei.demo.tftp.TftpData;
import top.ieei.demo.tftp.TftpError;
import top.ieei.demo.tftp.TftpException;
import top.ieei.demo.tftp.TftpOptionAck;
import top.ieei.demo.tftp.TftpPacket;
import top.ieei.demo.tftp.TftpRequest;
import top.ieei.demo.tftp.TftpSocket;


/**
 * 实际处理TFTP传输 <br>
 * @author xiecongshu
 */
public class ClientHandler
{
    /**
     * 日志
     */
    private static final Log LOG = LogFactory.getLog(ClientHandler.class);
    
    /**
     * 标记客户端连接是否已满
     */
    private boolean stuffedLink = false;
    
    /**
     * 客户端超时时间，毫秒
     */
    private int timeout = 5000;
    
    /**
     * 处理TFTP的读写请求的套接字封装对象
     */
    private TftpSocket tftpSocket;
    
    /**
     * 标记传输是否完成，完成的情况下，可以被TftpWorkPool回收资源，重新分配传输任务
     */
    private boolean isTransferComplete = true;
    
    /**
     * <默认构造函数>
     */
    public ClientHandler() throws SocketException
    {
        tftpSocket = new TftpSocket(timeout);
    }
    
    /**
     * 连接到客户端
     * @param inetAddress
     * @param clientPort void
     */
    public void connect(InetAddress inetAddress, int clientPort)
    {
        tftpSocket.connect(inetAddress, clientPort);
    }
    
    /**
     * 
     * 断开与客户端的连接
     * 
     */
    public void disconnect()
    {
        tftpSocket.disconnect();
    }
    
    /**
     * 获取当前会话的客户端的IP和端口信息
     */
    public String getClientInfo()
    {
        StringBuffer sb = new StringBuffer("TFTP Client:");
        String ip = "";
        int port = 0;
        if (null != tftpSocket)
        {
            InetAddress addr = tftpSocket.getAddress();
            if (null != addr)
            {
                ip = addr.getHostAddress();
                sb.append(ip);
            }
            port = tftpSocket.getPort();
            if (port != 0)
            {
                sb.append(":");
                sb.append(port);
            }
            sb.append(" ");
        }
        return sb.toString();
    }
    
    /**
     * 验证接收的数据包是否和预期的一致 
     * @param expecting 期望的ACK
     * @param received 接收到的ACK
     * @return true 如果匹配预期的结果
     */
    private boolean verifyAnswerPacket(TftpAck expectation, TftpAck received)
    {
        if (received == null)
        {
            return false;
        }
        
        if (expectation == null)
        {
            return false;
        }
        // 操作码不一致
        if (expectation.getOpCode() != received.getOpCode())
        {
            return false;
        }
        // 传输数据块编号不一致
        if (expectation.getBlockNum() != received.getBlockNum())
        {
            return false;
        }
        return true;
    }
    
    /**
     * 传输数据
     * @param send
     * @param recv
     * @return TftpAck
     */
    public TftpAck transferData(TftpAck send, TftpAck recv)
    {
        int retries = 0;
        int failCount = 0;
        
        // 发送数据包
        try
        {
            tftpSocket.write(send);
        }
        catch (Exception e)
        {
            LOG.debug(getClientInfo() + "TFTP send packet failure. "
                    + e.getMessage());
        }
        
        TftpPacket tftpPacket = null;
        boolean receiving = true;
        
        // 等待成功响应ACK
        while (receiving)
        {
            try
            {
                // 如果连接满了，设置两倍超时时间
                if (stuffedLink)
                {
                    tftpSocket.setSockTimeout(timeout * 2);
                }
                tftpPacket = tftpSocket.read();
            }
            catch (Exception e)
            {
                LOG.error("Exception: " + e.getMessage());
                return null;
            }
            
            // 如果没有收到响应
            if (tftpPacket == null)
            {
                if (retries++ > 5)
                {
                    // 重试次数太多，丢弃处理，发送错误响应给客户端
                    sendErrorPacket(TftpError.ERROR_NOT_DEFINED,
                            "Retransmit limit exceeded.");
                    LOG.debug(getClientInfo()
                            + " Maximum retransmit count exceeded.");
                    return null;
                }
                else
                {
                    // 重发数据并等待
                    LOG.debug(getClientInfo()
                            + "expected packet before time out, sending ACK/DATA again.");
                    try
                    {
                        tftpSocket.write(send);
                        
                        // 客户端数据处理管道可能已经充满，需要设置较长超时时间给客户端清理，并给出回应
                        stuffedLink = true;
                    }
                    catch (Exception e)
                    {
                        LOG.error("Send packet to " + getClientInfo()
                                + "failure. " + e.getMessage());
                    }
                    continue;
                }
            }
            
            // 收到出错的数据包，结束传输请求
            if (tftpPacket instanceof TftpError)
            {
                LOG.debug(getClientInfo()
                        + ((TftpError) tftpPacket).getErrorMessage());
                return null;
            }
            
            // 收到预期的回复响应
            if ((tftpPacket instanceof TftpAck)
                    && verifyAnswerPacket(recv, (TftpAck) tftpPacket))
            {
                try
                {
                    // 正确收到响应包 将超时时间改回默认
                    if (stuffedLink)
                    {
                        stuffedLink = false;
                        tftpSocket.setSockTimeout(timeout);
                    }
                }
                catch (SocketException e)
                {
                    LOG.debug(getClientInfo()
                            + "Could not change timeout on socket. "
                            + e.getMessage());
                }
                return (TftpAck) tftpPacket;
            }
            
            // 收到一些垃圾包，数量过多时，放弃
            if (failCount++ > 5)
            {
                return null;
            }
        }
        return null;
    }
    
    /**
     * 上传文件
     * @param os
     * @param clientAddress
     * @param clientPort
     * @param request
     * @return boolean
     */
    public boolean uploadFile(OutputStream os, InetAddress clientAddress,
            int clientPort, TftpRequest request)
    {
        setTransferComplete(false);
        // 连接客户端
        tftpSocket.connect(clientAddress, clientPort);
        
        // 当接收到写请请求时，发送第一个响应包ACK
        int blockNum = 0;
        
        // 文件数据
        byte[] data = new byte[1];
        
        // 已接收的文件大小
        int totalSize = 0;
        
        // 数据库大小
        int blkSize = 512;
        
        // 发送的响应包，实际可能是ACK或OACK
        TftpAck send;
        
        // 接收到的数据包
        TftpData receive;
        
        // 期望接收到的数据包
        TftpData expect;
        
        //传输文件的总大小
        int fileTotalSize = 0;
        
        // 如果请求数据包中包含选项，那么我们回应OACK而不是ACK
        if (request.hasOptions())
        {
            Options options = new Options();
            
            // 如果有块大小选项，且大于512，回应接受
            if (request.getBlkSize() > 512)
            {
                options.setBlockSize(request.getBlkSize());
                blkSize = request.getBlkSize();
            }
            
            // 如果有超时时间选项，回应系统默认超时时间
            if (request.getTimeout() > 0)
            {
                options.setTimeout(getTimeout());
            }
            
            // 如果有传输文件总大小选项，回应接受
            if (request.getTransferSize() > 0)
            {
                fileTotalSize = request.getTransferSize();
                options.setTransferSize(fileTotalSize);
            }
            send = new TftpOptionAck(0, options);
        }
        else
        {
            // 没有选项请求，回应上传请求接受
            send = new TftpAck(blockNum);
        }
        do
        {
            try
            {
                // 期望回复的包
                expect = new TftpData(blockNum + 1, data);
                
                // 发送第一个响应包
                receive = (TftpData) transferData(send, expect);
            }
            catch (TftpException e)
            {
                sendErrorPacket(TftpError.ERROR_NOT_DEFINED, e.getMessage());
                LOG.debug("TftpException:" + getClientInfo()
                        + e.getErrorMessage());
                setTransferComplete(true);
                return false;
            }
            catch (Exception e)
            {
                LOG.warn("Some unrecoverable error occured. cannot continue with client: "
                        + e.getMessage());
                setTransferComplete(true);
                return false;
            }
            
            // 如果 没有收到响应包 返回上传失败
            if (receive == null)
            {
                setTransferComplete(true);
                return false;
            }
            
            // 接收到上传的文件数据
            data = receive.getData();
            
            // 修改已经接收的文件的大小
            totalSize += data.length;
            
            LOG.debug(getClientInfo() + "received total file size is: "
                    + totalSize);
            
            // 检查接收到的数据块是否大于协商的数据库大小，大于协商的大小，返回错误信息
            if (data.length > blkSize)
            {
                sendErrorPacket(TftpError.ERROR_NOT_DEFINED,
                        "Exceed block size:" + blkSize);
                LOG.info(getClientInfo()
                        + ",Client tries to write file and exceed block size!");
                setTransferComplete(true);
                return false;
            }
            
            // 将接收到的数据写入文件
            try
            {
                os.write(data);
            }
            catch (IOException e)
            {
                sendErrorPacket(TftpError.ERROR_ILLEGAL_OP, e.getMessage());
                LOG.error("OutputStream write data error in uploadFile: "
                        + e.getMessage() + " caused by: "
                        + e.getCause().getMessage());
                setTransferComplete(true);
                return false;
            }
            blockNum = receive.getBlockNum();
            send = new TftpAck(blockNum);
        } while ((data.length == TftpSocket.BLOCK_SIZE)
                || (request.hasOptions() && request.getBlkSize() == data.length));
        
        // 在发送最后的数据确认包之前，关闭文件输出流
        try
        {
            os.close();
        }
        catch (IOException e)
        {
            LOG.error("problem while closing OutputStream in receiveFileFromClient: "
                    + e.getMessage()
                    + " caused by: "
                    + e.getCause().getMessage());
            setTransferComplete(true);
            return false;
        }
        
        // 发送最后的接收数据的确认包
        sendLastAckForFileUpload(receive);
        setTransferComplete(true);
        return true;
    }
    
    /**
     * 下载文件
     * @param is
     * @param clientAddress
     * @param clientPort
     * @param request
     * @return boolean
     */
    public boolean downloadFile(InputStream is, InetAddress clientAddress,
            int clientPort, TftpRequest request)
    {
        // 标记传输未完成
        setTransferComplete(false);
        
        // 检查请求参数
        if (clientAddress == null || clientPort == 0)
        {
            setTransferComplete(true);
            return false;
        }
        
        // 连接客户端
        tftpSocket.connect(clientAddress, clientPort);
        
        LOG.debug("Clearing socket buffer.");
        tftpSocket.clear();
        
        // 下载文件数据，数据缓冲
        byte[] data = new byte[TftpSocket.BLOCK_SIZE];
        
        if (null == is || null == request)
        {
            setTransferComplete(true);
            return false;
        }
        InputStream bis = new BufferedInputStream(is);
        
        // 读取数据长度
        int read = 0;
        
        // 数据块编号
        int blockNumber = 0;
        
        //  发送的响应包 ACK，OACK，DATD
        TftpAck send;
        
        // 期望的响应
        TftpAck expect;
        
        // 接收的包，ACK，DATD
        TftpAck receive;
        
        try
        {
            // 如果请求中包含选项信息，那么发送选项确认OACK
            if (request.hasOptions())
            {
                int length = is.available();
                Options options = new Options();
                if (request.getBlkSize() > 512)
                {
                    options.setBlockSize(request.getBlkSize());
                    data = new byte[request.getBlkSize()];
                }
                
                if (request.getTimeout() > 0)
                {
                    // 设置默认超时时间
                    options.setTimeout(getTimeout());
                }
                // 默认添加文件大小选项
                options.setTransferSize(length);
                send = new TftpOptionAck(0, options);
            }
            else
            {
                // 读取blockSize大小的数据到data中
                read = bis.read(data);
                
                // 第一次读取，如果读取数据小于或等于0，说明读取失败
                if (read <= 0)
                {
                    setTransferComplete(true);
                    return false;
                }
                send = new TftpData(++blockNumber, data, 0, read);
            }
            expect = new TftpAck(blockNumber);
            
            // 发送数据OACK或DATA包，并获取响应包
            receive = (TftpAck) transferData(send, expect);
            
            if (receive == null)
            {
                setTransferComplete(true);
                return false;
            }
            
            // 如果没有选项值，当实际读取数据大小小于512时，传输完成
            if (!request.hasOptions() && (read < TftpSocket.BLOCK_SIZE))
            {
                setTransferComplete(true);
                return true;
            }
        }
        catch (TftpException e)
        {
            LOG.error("Create DATA and ACK packet fail. " + e.getErrorMessage());
            sendErrorPacket(TftpError.ERROR_NOT_DEFINED, e.getMessage());
            setTransferComplete(true);
            return false;
        }
        catch (IOException e)
        {
            sendErrorPacket(TftpError.ERROR_NOT_DEFINED, e.getMessage());
            LOG.error(getClientInfo() + " " + e.getMessage());
            setTransferComplete(true);
            return false;
        }
        catch (Exception e)
        {
            LOG.error("Some unrecoverable error occured. cannot continue with client: "
                    + e.getMessage());
            setTransferComplete(true);
            return false;
        }
        
        do
        {
            try
            {
                read = bis.read(data);
                // 当读取数据EOF时，即读取完毕时，发送0数据块，表示最后下载数据完成
                if (read == -1)
                {
                    read = 0;
                }
                // 发送的数据包
                send = new TftpData(++blockNumber, data, 0, read);
                
                // 期望得到回应的数据包
                expect = new TftpAck(blockNumber);
                
                // 接收到的数据包
                receive = (TftpAck) transferData(send, expect);
                
                if (receive == null)
                {
                    setTransferComplete(true);
                    return false;
                }
            }
            catch (TftpException e)
            {
                LOG.error("Create DATA and ACK packet fail. "
                        + e.getErrorMessage());
                sendErrorPacket(TftpError.ERROR_NOT_DEFINED, e.getMessage());
                setTransferComplete(true);
                return false;
            }
            catch (IOException e)
            {
                sendErrorPacket(TftpError.ERROR_NOT_DEFINED, e.getMessage());
                LOG.error(getClientInfo() + e.getMessage());
                setTransferComplete(true);
                return false;
            }
            catch (Exception e)
            {
                LOG.error("Some unrecoverable error occured. cannot continue with client: "
                        + e.getMessage());
                setTransferComplete(true);
                return false;
            }
        } while ((read == TftpSocket.BLOCK_SIZE)
                || (request.hasOptions() && request.getBlkSize() == read));
        setTransferComplete(true);
        return true;
    }
    
    /**
     * 发送响应包
     * @param blocknumber void
     */
    private void sendAckPacket(int blockNumber)
    {
        try
        {
            TftpPacket tftpP = new TftpAck(blockNumber);
            tftpSocket.write(tftpP);
        }
        catch (Exception e)
        {
            LOG.error("Send TftpAck packet to:" + getClientInfo() + "failure. "
                    + e.getMessage());
        }
        LOG.debug("Send TftpAck packet to:" + getClientInfo() + "Block Number="
                + blockNumber);
    }
    
    /**
     * 发送出错包回应客户端
     * @param errorCode
     * @param errorMsg void
     */
    public void sendErrorPacket(int errorCode, String errorMsg)
    {
        try
        {
            TftpPacket tftpP = new TftpError(errorCode, errorMsg);
            tftpSocket.write(tftpP);
        }
        catch (Exception e)
        {
            LOG.error("Send TftpError packet to:" + getClientInfo()
                    + "failure. " + e.getMessage());
            return;
        }
        finally
        {
            setTransferComplete(true);
        }
        LOG.debug("Send TftpError packet to:" + getClientInfo() + "Error Info="
                + errorMsg);
    }
    
    /**
     * 获取响应超时时间，毫秒
     * 
     * @return timeout 毫秒
     */
    public int getTimeout()
    {
        return timeout;
    }
    
    /**
     * 设置响应超时时间，毫秒
     * 
     * @param timeout 超时时间，毫秒
     */
    public void setTimeout(int timeout)
    {
        this.timeout = timeout;
        try
        {
            if (tftpSocket != null)
                tftpSocket.setSockTimeout(timeout);
        }
        catch (SocketException e)
        {
            LOG.error("Could not set socket timeout on worker socket!");
        }
    }
    
    /**
     * 是否等待新的请求
     * @return boolean
     * @see [类、类#方法、类#成员]
     */
    public boolean waitingForNewRequest()
    {
        return isTransferComplete();
    }
    
    /**
     * 发送最后的响应包，不需要重发数据包
     */
    public void sendLastAckForFileUpload(TftpData data)
    {
        try
        {
            setTransferComplete(false);
            sendAckPacket(data.getBlockNum());
            
            // 等待一段时间，超时后重发确认信息
            tftpSocket.setSockTimeout(500);
            int count = timeout / 1000;
            while ((!isTransferComplete()) && (count > 0))
            {
                count--;
                TftpPacket tftpP = tftpSocket.read();
                if (tftpP == null)
                    continue;
                if (tftpP instanceof TftpData)
                {
                    sendAckPacket(data.getBlockNum());
                }
                if (tftpP instanceof TftpError)
                {
                    // 理论上会出现接收到错误信息的包，此时我们记录错误日志，实际上文件上传已经成功
                }
            }
        }
        catch (Exception e)
        {
            LOG.info("Could not wait to check if last ack was received by client!"
                    + e.getMessage());
        }
        finally
        {
            // 等待时间过去 客户端可以检测到回应的数据包
            setTransferComplete(true);
            setTimeout(timeout);
        }
    }
    
    /**
     * 设置传输是否完成
     * @param isComplete void
     */
    public synchronized void setTransferComplete(boolean isComplete)
    {
        this.isTransferComplete = isComplete;
    }
    
    /**
     * 获取传输是否完成
     * @return boolean
     */
    public synchronized boolean isTransferComplete()
    {
        return isTransferComplete;
    }
}
